from datetime import datetime
import cv2
import pybgs as bgs


def timeit(method):
    """装饰器：测量函数执行时间"""

    def timed(*args, **kw):
        start_time = datetime.now()
        result = method(*args, **kw)
        end_time = datetime.now()
        elapsed_time_ms = (end_time - start_time).total_seconds() * 1000  # 转换为毫秒
        print(f"{method.__name__} 耗时: {elapsed_time_ms:.0f} 毫秒")
        return result

    return timed


def initialize_video_capture(source):
    """初始化视频捕获对象"""
    return cv2.VideoCapture(source)


def get_video_properties(camera):
    """读取视频属性"""
    width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(camera.get(cv2.CAP_PROP_FPS))
    frame_count = int(camera.get(cv2.CAP_PROP_FRAME_COUNT))
    return width, height, fps, frame_count


# @timeit
def process_frame(frame, algorithm, rectangle, counter, pre_frame, threshold):
    """处理单个帧"""
    rectangleX, rectangleY, rectangleXCols, rectangleYCols = rectangle
    cv2.rectangle(frame, (rectangleX, rectangleY),
                  (rectangleX + rectangleXCols, rectangleY + rectangleYCols), (0, 255, 0), 2)

    if pre_frame is not None:
        img_delta = algorithm.apply(frame)
        roi = img_delta[rectangleY:rectangleY + rectangleYCols, rectangleX:rectangleX + rectangleXCols]
        thresh = cv2.threshold(roi, 25, 255, cv2.THRESH_BINARY)[1]
        thresh = cv2.dilate(thresh, None, iterations=2)
        contours, hierarchy = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        area = sum(cv2.contourArea(x) for x in contours)
        warn_color = (0, 0, 255) if area > threshold else (55, 255, 155)
        cv2.putText(frame, 'warn index  %s' % area, (rectangleX, rectangleY), cv2.FONT_HERSHEY_SIMPLEX, 2, warn_color,
                    2)
        return roi, frame, area

    return None, frame, 0


def save_and_display(frame, writer,winname="Processed Frame"):
    """保存和显示结果"""
    writer.write(frame)
    cv2.imshow(winname, frame)
    cv2.waitKey(1)


def release_resources(camera, writers):
    """释放资源"""
    camera.release()
    [writer.release() for writer in writers]
    cv2.destroyAllWindows()


@timeit
def main_process(file, fname):
    camera = initialize_video_capture(file)
    width, height, fps, counterNum = get_video_properties(camera)
    rectangle = (450, 260, 800, 500)  # 监测区域坐标和大小
    threshold = 100  # 阈值
    algorithm = bgs.ViBe()  # 背景减法算法
    videowriter_frame = cv2.VideoWriter(file.replace('.', '_result.'),
                                  cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
    videowriter_fg = cv2.VideoWriter(file.replace('.', '_result_fg.'),
                                        cv2.VideoWriter_fourcc(*"mp4v"), fps, (rectangle[2], rectangle[3]))
    pre_frame = None
    counter = 1
    while True:
        read_start_time = datetime.now()
        grabbed, frame = camera.read()
        read_time = datetime.now() - read_start_time

        if not grabbed:
            break

        if pre_frame is None:
            pre_frame = frame


        if counter % (fps // fps) == 0:
            process_start_time = datetime.now()

            # print(frame)
            pre_frame, frame, area = process_frame(frame, algorithm, rectangle, counter, pre_frame,
                                                   threshold)
            process_time = datetime.now() - process_start_time

            write_start_time = datetime.now()
            save_and_display(frame, videowriter_frame,"frame")
            foreground_c= cv2.cvtColor(pre_frame, cv2.COLOR_GRAY2BGR)
            save_and_display(foreground_c, videowriter_fg,"foreground")
            write_time = datetime.now() - write_start_time

            print(
                f"处理第{counter}帧/共{counterNum}帧,  "
                f"读取耗时: {read_time.total_seconds() * 1000:.0f} 毫秒, "
                f"处理耗时: {process_time.total_seconds() * 1000:.0f} 毫秒, "
                f"写入耗时: {write_time.total_seconds() * 1000:.0f} 毫秒"
            )

        counter += 1
    release_resources(camera, [videowriter_frame,videowriter_fg])


if __name__ == "__main__":
    file = "data/suidao_stable.mp4"
    fname = "suidao"
    # user, pwd, ip, channel = "admin", "buaa0707", "192.168.1.103", 1  # 主码流
    # # user, pwd, ip, channel = "admin", "buaa0707", "192.168.1.103", 2  # 子码流 分辨率低
    # video_stream_path = "rtsp://%s:%s@%s//Streaming/Channels/%d" % (user, pwd, ip, channel)  # HIKIVISION new version 2017
    main_process(file, fname)
